import torch
from torch.utils.data import Dataset
from torchvision.transforms import transforms


# 处理数据集程序
class CustomDataset(Dataset):
    def __init__(self, data, noisy=False):
        self.data = data
        self.noisy = noisy

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        x = torch.tensor(self.data[index, :-1], dtype=torch.float32)
        if self.noisy:
            mean = 0
            std = 0.1
            noise = torch.normal(mean, std, x.size())
            x = x + noise
        y = torch.tensor(self.data[index, -1], dtype=torch.long)
        return x, y


class AddGaussianNoise(object):
    def __init__(self, mean=0, std=0.1):
        self.mean = mean
        self.std = std

    def __call__(self, tensor):
        return tensor + torch.randn_like(tensor) * self.std + self.mean
